from flask import Flask, render_template, request, redirect, url_for, flash, jsonify, session, send_file, make_response
import pandas as pd
import os
import json
import numpy as np
from datetime import datetime
import hashlib
import secrets
import io
from functools import wraps

# 导入预测模型相关函数
from test import predict_transaction_errors
from models import *
# 导入数据导入功能
from import_data import import_data_from_csv

# 创建Flask应用
app = Flask(__name__)
app.secret_key = secrets.token_hex(16)  # 生成随机密钥

# 添加自定义Jinja2过滤器
@app.template_filter('timestamp_to_date')
def timestamp_to_date(timestamp):
    """将时间戳转换为可读的日期时间格式"""
    return datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S')

# 添加max和min过滤器供模板使用
@app.template_filter('max_value')
def max_value(a, b):
    """返回两个值中的较大值"""
    return max(a, b)

@app.template_filter('min_value')
def min_value(a, b):
    """返回两个值中的较小值"""
    return min(a, b)


# 登录验证装饰器
def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user_id' not in session:
            flash('请先登录', 'danger')
            return redirect(url_for('login'))
        return f(*args, **kwargs)
    return decorated_function

# 首页路由
@app.route('/')
def index():
    # 获取基本统计数据供首页展示
    try:
        total_transactions = Transaction.select().count()
        
        # 获取唯一地址数
        from_addresses = Transaction.select(Transaction.from_address).distinct()
        to_addresses = Transaction.select(Transaction.to_address).distinct()
        unique_addresses = set([tx.from_address for tx in from_addresses] + [tx.to_address for tx in to_addresses])
        total_addresses = len(unique_addresses)
        
        # 获取最近5笔交易
        recent_transactions = Transaction.select().order_by(Transaction.timestamp.desc()).limit(5)
        
        return render_template('index.html', 
                              total_transactions=total_transactions,
                              total_addresses=total_addresses,
                              recent_transactions=recent_transactions)
    except Exception as e:
        # 如果出错，仍然渲染首页，但不传递数据
        return render_template('index.html')

# 用户注册路由
@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        email = request.form.get('email')
        
        # 验证用户名是否已存在
        try:
            User.get(User.username == username)
            flash('用户名已存在', 'danger')
            return render_template('register.html')
        except User.DoesNotExist:
            pass
        
        # 创建新用户
        hashed_password = hashlib.sha256(password.encode()).hexdigest()
        User.create(
            username=username,
            password=hashed_password,
            email=email
        )
        
        flash('注册成功，请登录', 'success')
        return redirect(url_for('login'))
    
    return render_template('register.html')

# 用户登录路由
@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        
        try:
            user = User.get(User.username == username)
            hashed_password = hashlib.sha256(password.encode()).hexdigest()
            
            if user.password == hashed_password:
                # 更新最后登录时间
                user.last_login = datetime.now()
                user.save()
                
                # 设置会话
                session['user_id'] = user.id
                session['username'] = user.username
                
                flash('登录成功', 'success')
                return redirect(url_for('dashboard'))
            else:
                flash('密码错误', 'danger')
        except User.DoesNotExist:
            flash('用户不存在', 'danger')
    
    return render_template('login.html')

# 用户登出路由
@app.route('/logout')
def logout():
    session.clear()
    flash('已退出登录', 'info')
    return redirect(url_for('index'))

# 用户仪表板
@app.route('/dashboard')
@login_required
def dashboard():
    # 获取交易数据统计信息
    try:
        total_transactions = Transaction.select().count()
        unique_addresses = set()
        
        for tx in Transaction.select(Transaction.from_address, Transaction.to_address):
            unique_addresses.add(tx.from_address)
            unique_addresses.add(tx.to_address)
        
        total_addresses = len(unique_addresses)
        
        # 获取最近10笔交易
        recent_transactions = Transaction.select().order_by(Transaction.timestamp.desc()).limit(10)
        
        return render_template('dashboard.html', 
                              total_transactions=total_transactions,
                              total_addresses=total_addresses,
                              recent_transactions=recent_transactions)
    except Exception as e:
        flash(f'获取数据时出错: {str(e)}', 'danger')
        return render_template('dashboard.html')

# 个人信息页面
@app.route('/profile', methods=['GET', 'POST'])
@login_required
def profile():
    user = User.get(User.id == session['user_id'])
    
    if request.method == 'POST':
        email = request.form.get('email')
        
        # 更新用户信息
        user.email = email
        user.save()
        
        flash('个人信息已更新', 'success')
    
    return render_template('profile.html', user=user)

# 修改密码
@app.route('/change_password', methods=['POST'])
@login_required
def change_password():
    user = User.get(User.id == session['user_id'])
    current_password = request.form.get('current_password')
    new_password = request.form.get('new_password')
    confirm_password = request.form.get('confirm_password')
    
    # 验证当前密码
    hashed_current = hashlib.sha256(current_password.encode()).hexdigest()
    if user.password != hashed_current:
        flash('当前密码不正确', 'danger')
        return redirect(url_for('profile'))
    
    # 验证新密码
    if new_password != confirm_password:
        flash('新密码与确认密码不匹配', 'danger')
        return redirect(url_for('profile'))
    
    # 更新密码
    hashed_new = hashlib.sha256(new_password.encode()).hexdigest()
    user.password = hashed_new
    user.save()
    
    flash('密码已更新', 'success')
    return redirect(url_for('profile'))

# 交易数据列表
@app.route('/transactions')
@login_required
def transactions():
    page = request.args.get('page', 1, type=int)
    per_page = 20
    
    # 获取搜索参数
    address = request.args.get('address', '')
    min_value = request.args.get('min_value', '')
    max_value = request.args.get('max_value', '')
    
    # 构建查询
    query = Transaction.select()
    
    # 应用过滤条件
    if address:
        query = query.where(
            (Transaction.from_address.contains(address)) | 
            (Transaction.to_address.contains(address))
        )
    
    if min_value and min_value.replace('.', '', 1).isdigit():
        query = query.where(Transaction.value >= float(min_value))
    
    if max_value and max_value.replace('.', '', 1).isdigit():
        query = query.where(Transaction.value <= float(max_value))
    
    # 计算总数和分页
    total_count = query.count()
    total_pages = (total_count + per_page - 1) // per_page
    
    # 应用排序和分页
    transactions = query.order_by(Transaction.timestamp.desc()).paginate(page, per_page)
    
    return render_template('transactions.html', 
                          transactions=transactions,
                          page=page,
                          total_pages=total_pages)

# 交易数据汇总
@app.route('/transactions/summary')
@login_required
def transaction_summary():
    # 获取交易数据汇总信息
    try:
        # 总交易数
        total_transactions = Transaction.select().count()
        
        # 总交易金额
        total_value = Transaction.select(pw.fn.SUM(Transaction.value)).scalar()
        
        # 唯一地址数
        from_addresses = Transaction.select(Transaction.from_address).distinct()
        to_addresses = Transaction.select(Transaction.to_address).distinct()
        unique_addresses = set([tx.from_address for tx in from_addresses] + [tx.to_address for tx in to_addresses])
        total_addresses = len(unique_addresses)
        
        # 按时间段统计交易数量
        # 获取最早和最晚的时间戳
        earliest_tx = Transaction.select(pw.fn.MIN(Transaction.timestamp)).scalar()
        latest_tx = Transaction.select(pw.fn.MAX(Transaction.timestamp)).scalar()
        
        # 计算时间范围并分段
        time_range = latest_tx - earliest_tx
        segments = 10  # 分成10个时间段
        segment_size = time_range / segments
        
        time_distribution = []
        for i in range(segments):
            start_time = earliest_tx + i * segment_size
            end_time = earliest_tx + (i + 1) * segment_size
            count = Transaction.select().where(
                (Transaction.timestamp >= start_time) & 
                (Transaction.timestamp < end_time)
            ).count()
            time_distribution.append({
                'segment': i + 1,
                'start_time': datetime.fromtimestamp(start_time).strftime('%Y-%m-%d'),
                'end_time': datetime.fromtimestamp(end_time).strftime('%Y-%m-%d'),
                'count': count
            })
        
        return render_template('transaction_summary.html',
                              total_transactions=total_transactions,
                              total_value=total_value,
                              total_addresses=total_addresses,
                              time_distribution=time_distribution)
    except Exception as e:
        flash(f'获取汇总数据时出错: {str(e)}', 'danger')
        return render_template('transaction_summary.html')

# 导出全部交易数据为CSV
@app.route('/transactions/export')
@login_required
def export_transactions():
    try:
        # 查询所有交易数据
        transactions = Transaction.select()
        
        # 创建DataFrame
        data = [{
            'TxHash': tx.tx_hash,
            'BlockHeight': tx.block_height,
            'TimeStamp': tx.timestamp,
            'From': tx.from_address,
            'To': tx.to_address,
            'Value': tx.value,
            'IsError': 1 if tx.is_error else 0
        } for tx in transactions]
        
        df = pd.DataFrame(data)
        
        # 创建内存文件对象
        output = io.StringIO()
        
        # 导出为CSV到内存
        df.to_csv(output, index=False)
        
        # 将指针移到开头
        output.seek(0)
        
        # 生成文件名
        timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
        filename = f'transactions_{timestamp}.csv'
        
        # 创建响应对象
        response = make_response(output.getvalue())
        response.headers["Content-Disposition"] = f"attachment; filename={filename}"
        response.headers["Content-type"] = "text/csv"
        
        flash(f'正在下载交易数据', 'success')
        return response
    except Exception as e:
        flash(f'导出数据时出错: {str(e)}', 'danger')
        return redirect(url_for('transaction_summary'))

# 导出分析后的汇总数据为CSV
@app.route('/transactions/export_summary')
@login_required
def export_transactions_summary():
    try:
        # 获取交易数据汇总信息
        # 总交易数
        total_transactions = Transaction.select().count()
        
        # 总交易金额
        total_value = Transaction.select(pw.fn.SUM(Transaction.value)).scalar()
        
        # 唯一地址数
        from_addresses = Transaction.select(Transaction.from_address).distinct()
        to_addresses = Transaction.select(Transaction.to_address).distinct()
        unique_addresses = set([tx.from_address for tx in from_addresses] + [tx.to_address for tx in to_addresses])
        total_addresses = len(unique_addresses)
        
        # 按时间段统计交易数量
        earliest_tx = Transaction.select(pw.fn.MIN(Transaction.timestamp)).scalar()
        latest_tx = Transaction.select(pw.fn.MAX(Transaction.timestamp)).scalar()
        
        # 计算时间范围并分段
        time_range = latest_tx - earliest_tx
        segments = 10  # 分成10个时间段
        segment_size = time_range / segments
        
        time_distribution = []
        for i in range(segments):
            start_time = earliest_tx + i * segment_size
            end_time = earliest_tx + (i + 1) * segment_size
            count = Transaction.select().where(
                (Transaction.timestamp >= start_time) & 
                (Transaction.timestamp < end_time)
            ).count()
            time_distribution.append({
                'segment': i + 1,
                'start_time': datetime.fromtimestamp(start_time).strftime('%Y-%m-%d'),
                'end_time': datetime.fromtimestamp(end_time).strftime('%Y-%m-%d'),
                'count': count,
                'percentage': round(count / total_transactions * 100, 2) if total_transactions > 0 else 0
            })
        
        # 创建汇总数据DataFrame
        summary_data = {
            '总体统计': [{
                '总交易数': total_transactions,
                '总交易金额': float(total_value) if total_value else 0,
                '总地址数': total_addresses
            }],
            '时间分布': time_distribution
        }
        
        # 创建内存文件对象
        output = io.StringIO()
        
        # 生成文件名
        timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
        filename = f'transactions_summary_{timestamp}.csv'
        
        # 写入CSV数据到内存
        output.write('交易数据汇总报告\n\n')
        output.write('总体统计:\n')
        output.write(f'总交易数,{total_transactions}\n')
        output.write(f'总交易金额,{float(total_value) if total_value else 0}\n')
        output.write(f'总地址数,{total_addresses}\n\n')
        
        output.write('时间分布:\n')
        output.write('时间段,开始日期,结束日期,交易数量,占比(%)\n')
        for item in time_distribution:
            output.write(f"{item['segment']},{item['start_time']},{item['end_time']},{item['count']},{item['percentage']}\n")
        
        # 将指针移到开头
        output.seek(0)
        
        # 创建响应对象
        response = make_response(output.getvalue())
        response.headers["Content-Disposition"] = f"attachment; filename={filename}"
        response.headers["Content-type"] = "text/csv"
        
        flash(f'正在下载汇总分析数据', 'success')
        return response
    except Exception as e:
        flash(f'导出汇总数据时出错: {str(e)}', 'danger')
        return redirect(url_for('transaction_summary'))

# 数据可视化页面
@app.route('/visualization')
@login_required
def visualization():
    return render_template('visualization.html')

# 获取可视化数据的API
@app.route('/api/visualization/data')
@login_required
def visualization_data():
    try:
        # 1. 交易金额分布
        value_ranges = [
            {'min': 0, 'max': 0.1},
            {'min': 0.1, 'max': 1},
            {'min': 1, 'max': 10},
            {'min': 10, 'max': 100},
            {'min': 100, 'max': float('inf')}
        ]
        
        value_distribution = []
        for i, range_info in enumerate(value_ranges):
            min_val = range_info['min']
            max_val = range_info['max']
            
            if max_val == float('inf'):
                count = Transaction.select().where(Transaction.value >= min_val).count()
                label = f'>{min_val}'
            else:
                count = Transaction.select().where(
                    (Transaction.value >= min_val) & (Transaction.value < max_val)
                ).count()
                label = f'{min_val}-{max_val}'
            
            value_distribution.append({
                'range': label,
                'count': count
            })
        
        # 2. 用户交易次数占比（取前10个活跃地址）
        from_counts = {}
        for tx in Transaction.select(Transaction.from_address):
            from_counts[tx.from_address] = from_counts.get(tx.from_address, 0) + 1
        
        # 排序并取前10个
        top_addresses = sorted(from_counts.items(), key=lambda x: x[1], reverse=True)[:10]
        address_distribution = [{'address': addr, 'count': count} for addr, count in top_addresses]
        
        # 3. 用户近期交易金额趋势（取最近30天，按天汇总）
        latest_tx = Transaction.select(pw.fn.MAX(Transaction.timestamp)).scalar()
        thirty_days_ago = latest_tx - (30 * 24 * 60 * 60)  # 30天的秒数
        
        # 按天分组查询
        daily_values = []
        for day in range(30):
            day_start = thirty_days_ago + (day * 24 * 60 * 60)
            day_end = day_start + (24 * 60 * 60)
            
            daily_sum = Transaction.select(pw.fn.SUM(Transaction.value)).where(
                (Transaction.timestamp >= day_start) & (Transaction.timestamp < day_end)
            ).scalar() or 0
            
            daily_values.append({
                'day': datetime.fromtimestamp(day_start).strftime('%m-%d'),
                'value': float(daily_sum)
            })
        
        return jsonify({
            'value_distribution': value_distribution,
            'address_distribution': address_distribution,
            'daily_values': daily_values
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

# 非法账号识别页面
@app.route('/fraud_detection')
@login_required
def fraud_detection():
    return render_template('fraud_detection.html')

# 执行非法账号识别
@app.route('/api/fraud_detection/analyze', methods=['POST'])
@login_required
def analyze_fraud():
    try:
        # 获取要分析的地址和模型类型
        data = request.json
        address = data.get('address')
        model_type = data.get('model_type', 'XGBoost')  # 默认使用XGBoost模型
        
        # 验证模型类型
        valid_models = ['XGBoost', 'RandomForest', 'PDTGA', 'GAT', 'MLP']
        if model_type not in valid_models:
            return jsonify({'error': f'不支持的模型类型: {model_type}，支持的模型类型有: {valid_models}'}), 400
        
        if not address:
            # 如果没有指定地址，随机选择一些交易进行分析
            transactions = Transaction.select().order_by(pw.fn.Random()).limit(10)
        else:
            # 查询与该地址相关的交易
            transactions = Transaction.select().where(
                (Transaction.from_address == address) | (Transaction.to_address == address)
            ).limit(20)
        
        # 转换为预测模型需要的格式
        tx_data = [{
            'TxHash': tx.tx_hash,
            'BlockHeight': tx.block_height,
            'TimeStamp': tx.timestamp,
            'From': tx.from_address,
            'To': tx.to_address,
            'Value': tx.value
        } for tx in transactions]
        
        if not tx_data:
            return jsonify({'error': '没有找到相关交易数据'}), 404
        
        # 调用预测模型，传入模型类型
        results = predict_transaction_errors(tx_data, model_name=model_type)
        
        # 处理结果
        processed_results = []
        for addr, prob in results.items():
            processed_results.append({
                'address': addr,
                'probability': prob,
                'status': 'Likely Fraudulent' if prob > 0.5 else 'Likely Legitimate'
            })
        
        # 读取模型性能数据
        try:
            with open('model/model_performance.json', 'r') as f:
                model_performance_data = json.load(f)
                model_performance = model_performance_data.get(model_type, {})
        except Exception as e:
            print(f"读取模型性能数据出错: {e}")
            model_performance = {
                'accuracy': 0.85,
                'precision': 0.83,
                'recall': 0.82,
                'f1_score': 0.83
            }
        
        return jsonify({
            'results': processed_results,
            'transactions': tx_data,
            'model_type': model_type,  # 返回使用的模型类型
            'model_name': model_type,  # 添加模型名称
            'model_performance': model_performance  # 添加模型性能数据
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

# 数据导入页面
@app.route('/import_data')
@login_required
def import_data_page():
    return render_template('import_data.html')

# 数据导入API
@app.route('/api/import_data', methods=['POST'])
@login_required
def api_import_data():
    try:
        # 检查是否有文件上传
        if 'csv_file' not in request.files:
            return jsonify({'success': False, 'error': '没有上传文件'}), 400
        
        file = request.files['csv_file']
        
        # 检查文件名
        if file.filename == '':
            return jsonify({'success': False, 'error': '未选择文件'}), 400
        
        # 检查文件类型
        if not file.filename.endswith('.csv'):
            return jsonify({'success': False, 'error': '只支持CSV文件格式'}), 400
        
        # 保存上传的文件到临时目录
        temp_dir = 'temp'
        if not os.path.exists(temp_dir):
            os.makedirs(temp_dir)
        
        temp_path = os.path.join(temp_dir, f"upload_{datetime.now().strftime('%Y%m%d%H%M%S')}.csv")
        file.save(temp_path)
        
        # 调用导入函数
        result = import_data_from_csv(temp_path)
        
        # 删除临时文件
        try:
            os.remove(temp_path)
        except:
            pass
        
        # 处理导入结果
        if not result:
            return jsonify({
                'success': False, 
                'error': '导入过程中发生错误，请检查文件格式是否正确'
            }), 500
        
        # 返回成功结果
        total = result.get('imported', 0) + result.get('skipped', 0)
        return jsonify({
            'success': True,
            'total': total,
            'imported': result.get('imported', 0),
            'skipped': result.get('skipped', 0),
            'failed': 0
        })
        
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500

# 主函数
if __name__ == '__main__':
    # 创建数据库表
    create_tables()
    
    # 启动Flask应用
    app.run(debug=True, port=5000)